'''
对三分类数据集（腺癌，鳞癌，其他）进行分布统计
绘制统计盒图
针对性别特征，绘制统计直方图
'''



import matplotlib.pyplot as plt
import pandas as pd
import os

# 读取统计数据
def read_data():
    csv_path = 'D:/lung_cancer/data/divide_csv/three/train.csv'
    data = pd.read_csv(csv_path)
    return data

# 将某类数据按类别划分
def devide_data(data, cancer_type):

    x_list = [[], [], []]
    for i in range(len(data)):
        x_list[cancer_type[i]-1].append(float(data[i]))

    return x_list


def draw_box(one_data, name):
    item = name.split('/')[-1]
    print('------start draw ', item, '------')
    plt.figure(figsize=(8, 6))
    plt.boxplot(one_data, notch=False, sym='rs', vert=True)
    plt.xticks([y + 1 for y in range(3)], ['1', '2', '3'])
    plt.xlabel('measurement '+item)
    plt.title('Box plot '+item)
    # plt.savefig(name)
    plt.show()

# 对性别做统计，直方图更合适
def draw_hist(one_data, name):
    one_men = 0
    one_women = 0
    two_men = 0
    two_women = 0
    three_men = 0
    three_women = 0
    one_list = one_data[0]
    two_list = one_data[1]
    three_list = one_data[2]
    for item1 in one_list:
        if int(item1) == 0:
            one_men = one_men+1
        else:
            one_women = one_women+1
    for item2 in two_list:
        if int(item2) == 0:
            two_men = two_men+1
        else:
            two_women = two_women+1
    for item3 in three_list:
        if int(item3) == 0:
            three_men = three_men+1
        else:
            three_women = three_women+1

    name_list = ['1', '2', '3']
    x = list(range(len(name_list)))
    total_width, n = 0.8, 2
    width = total_width/n
    plt.bar(x, [one_men, two_men, three_men], width=width, label='men', fc='b')
    for i in range(len(x)):
        x[i] += width
    plt.bar(x, [one_women, two_women, three_women], width=width, label='women', tick_label=name_list, fc='g')
    plt.legend()
    # plt.savefig(name)
    plt.show()


# 对不同特征做盒图统计
def statistic():
    data = read_data()
    cancer_type = data['cancer_type']

    save_path = 'D:/lung_cancer/lung_cancer_code/release_code/output/three/Boxplot'
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    # 统计病灶半径
    cancer_r = data['r']
    r_list = devide_data(cancer_r, cancer_type)
    draw_box(r_list, save_path+'/r')

    # 统计病灶 X
    cancer_r = data['x']
    r_list = devide_data(cancer_r, cancer_type)
    draw_box(r_list, save_path + '/x')

    # 统计病灶 Y
    cancer_r = data['y']
    r_list = devide_data(cancer_r, cancer_type)
    draw_box(r_list, save_path + '/y')

    # suv_max
    suvmax = data['global_suvmax']
    suvmax_list = devide_data(suvmax, cancer_type)
    draw_box(suvmax_list, save_path+'/suv_max')

    # suv_min
    suvmin = data['local_suvmin']
    suvmin_list = devide_data(suvmin, cancer_type)
    draw_box(suvmin_list, save_path+'/suv_min')

    # suv_avg
    suvavg = data['local_suvavg']
    suvavg_list = devide_data(suvavg, cancer_type)
    draw_box(suvavg_list, save_path+'/suv_avg')

    # suv_std
    suvstd = data['local_suvstd']
    suvstd_list = devide_data(suvstd, cancer_type)
    draw_box(suvstd_list, save_path+'/suv_std')

    # suv_var
    suvstd = data['local_suvvar']
    suvstd_list = devide_data(suvstd, cancer_type)
    draw_box(suvstd_list, save_path + '/suv_var')

    # age
    age = data['patientAge']
    age_list = devide_data(age, cancer_type)
    draw_box(age_list, save_path+'/age')

    # size
    age = data['patientSize']
    age_list = devide_data(age, cancer_type)
    draw_box(age_list, save_path + '/size')

    # weight
    patientWeight = data['patientWeight']
    patientWeight_list = devide_data(patientWeight, cancer_type)
    draw_box(patientWeight_list, save_path+'/patientWeight')

    # Sex,性别用直方图统计更好
    patientSex = data['patientSex']
    patientSex_list = devide_data(patientSex, cancer_type)
    draw_hist(patientSex_list, save_path+'/patientSex')


if __name__ == '__main__':

    statistic()